﻿module Physics.Chapter5.Algebra

open Casanova.Core
open Casanova.Math
open Casanova.StandardLibrary.Core

type Matrix = { A : float32[][] }
  with
    static member Create(A) = { A = A }
    static member (*) (m1:Matrix, m2:Matrix) =
      if m1.A.Length = 0 then [||] |> Matrix.Create
      else
        let n,m,k = m1.A.Length, m2.A.Length, m2.A.[0].Length
        let A = 
          [|
            for i = 0 to n - 1 do
              yield
                [|
                  for l = 0 to k - 1 do
                    yield 0.0f
                |]
          |]
        for i = 0 to n - 1 do
          for l = 0 to k - 1 do
            for j = 0 to m - 1 do
              A.[i].[l] <- A.[i].[l] + m1.A.[i].[j] * m2.A.[j].[l]
        A |> Matrix.Create
    static member (*) (m:Matrix, v:Vector) =
      let n,l = m.A.Length, v.V.Length
      let V  = 
        [|
          for i = 0 to n - 1 do
            let sum = var 0.0f
            for j = 0 to l - 1 do 
              sum := !sum + m.A.[i].[j] * v.V.[j]
            yield !sum
        |] 
      V |> Vector.Create
    static member (*) (m:Matrix, f:float32<'a>) =
      if m.A.Length = 0 then [||] |> Matrix.Create
      else
        let n,l = m.A.Length, m.A.[0].Length
        [|
          for i = 0 to n - 1 do
            yield
              [|
                for j = 0 to l - 1 do
                  yield float32(m.A.[i].[j] * f)
              |]
        |] |> Matrix.Create
    member m.Transpose = 
      if m.A.Length = 0 then [||] |> Matrix.Create
      else
        let n,l = m.A.Length, m.A.[0].Length
        [|
          for j = 0 to l - 1 do
            yield
              [|
                for i = 0 to n - 1 do
                  yield m.A.[i].[j]
              |]
        |] |> Matrix.Create
    member m.ToArray2D =
      if m.A.Length = 0 then Array2D.zeroCreate 0 0
      else
        let n,l = m.A.Length, m.A.[0].Length
        Array2D.init n l (fun i j -> m.A.[i].[j])

and SparseMatrix = { A : (int * float32)[][]; NumRows : int; NumColumns : int }
  with
    static member Create(A : (int * float32)[][], num_rows : int, num_columns : int) = 
      { 
        A = A
        NumRows = num_rows
        NumColumns = num_columns
      }

and Vector = { V : float32[] }
  with 
    member v.Item with get i = v.V.[i]
    static member Create(V) = { V = V }
    static member (*) (v:Vector, f:float32<'a>) =
      let n = v.V.Length
      [|
        for i = 0 to n - 1 do
          yield float32(v.V.[i] * f)
      |] |> Vector.Create
    static member (*) (v1:Vector, v2:Vector) =
      let n = v1.V.Length
      seq{
        for i = 0 to n - 1 do
          yield float32(v1.V.[i] * v2.V.[i])
      } |> Seq.sum
    static member (/) (v:Vector, f:float32<'a>) = v * (1.0f / f)
    static member (+) (v1:Vector, v2:Vector) = 
      let n = v1.V.Length
      [|
        for i = 0 to n - 1 do
          yield v1.V.[i] + v2.V.[i]
      |] |> Vector.Create
    static member (-) (v1:Vector, v2:Vector) = 
      let n = v1.V.Length
      [|
        for i = 0 to n - 1 do
          yield v1.V.[i] - v2.V.[i]
      |] |> Vector.Create

let vec_to_seq (v:Vector3<'a>) = seq{ yield v.X |> float32; yield v.Y |> float32; yield v.Z |> float32 }

let clamp(x,x_min, x_max) = (max x_min >> min x_max) x

let projected_gauss_seidel (A:SparseMatrix) (b:Vector) (x0:Vector) (x_min:Vector) (x_max:Vector) =
  let x0,x_min,x_max = x0.V, x_min.V, x_max.V
  let x = [| for v in x0 do yield v |] |> Vector.Create
  let rec loop n = 
    let mutable tot_dx = 0.0f
    for i = 0 to x.V.Length - 1 do
      let mutable Ai_x = 0.0f
      let mutable a_ii = 0.0f
      for j,a_ij in A.A.[i] do 
        Ai_x <- Ai_x + x.V.[j] * a_ij
        if j = i then a_ii <- a_ij
      let dx = (b.[i] - Ai_x) / a_ii
      let x_i' = clamp(x.V.[i] + dx, x_min.[i], x_max.[i])
      do tot_dx <- tot_dx + abs(x_i' - x.V.[i])
      do x.V.[i] <- x_i'
    if n > 0 && tot_dx / (float32 x.V.Length) > 0.001f then
      do loop (n-1)
  do loop 100
  x
